package ipaX

import (
	"math"

	"github.com/crate-crypto/go-ipa/bandersnatch"
	"github.com/crate-crypto/go-ipa/bandersnatch/fr"
	"github.com/crate-crypto/go-ipa/common"
	"github.com/crate-crypto/go-ipa/ipa"
)

// 分成相等的俩个 slice
func splitScalars(x []fr.Element) ([]fr.Element, []fr.Element) {
	if len(x)%2 != 0 {
		panic("slice should have an even length")
	}

	mid := len(x) / 2
	return x[:mid], x[mid:]
}

// 分成相等的俩个 slice
func splitPoints(x []bandersnatch.PointAffine) ([]bandersnatch.PointAffine, []bandersnatch.PointAffine) {
	if len(x)%2 != 0 {
		panic("slice should have an even length")
	}

	mid := len(x) / 2
	return x[:mid], x[mid:]
}

// 生成多项式系数，256个
func polynomial256(polynomial ...uint64) []fr.Element {
	n := len(polynomial)
	if len(polynomial) > 256 {
		panic("polynomial cannot exceed 256 coefficients")
	}
	polynomialFr := make([]fr.Element, 256)
	for i := 0; i < n; i++ {
		polynomialFr[i].SetUint64(polynomial[i])
	}

	pad := 256 - n
	for i := n; i < pad; i++ {
		polynomialFr[i] = fr.Zero()
	}

	return polynomialFr
}

func compute_num_rounds(vector_size uint32) uint32 {
	// Check if this number is 0
	// zero is not a valid input to this function for our usecase
	if vector_size == 0 {
		panic("zero is not a valid input")
	}

	// See: https://stackoverflow.com/a/600306
	isPow2 := (vector_size & (vector_size - 1)) == 0

	if !isPow2 {
		panic("non power of 2 numbers are not valid inputs")
	}

	res := math.Log2(float64(vector_size))

	return uint32(res)
}

// Computes c[i] =a[i] + b[i] * x
// returns c
// panics if len(a) != len(b)
func foldScalars(a []fr.Element, b []fr.Element, x fr.Element) []fr.Element {

	if len(a) != len(b) {
		panic("slices not equal length")
	}

	result := make([]fr.Element, len(a))
	for i := 0; i < len(a); i++ {
		var bx fr.Element
		bx.Mul(&x, &b[i])
		result[i].Add(&bx, &a[i])
	}
	return result
}

// Computes c[i] =a[i] + b[i] * x
// returns c
// panics if len(a) != len(b)
func foldPoints(a []bandersnatch.PointAffine, b []bandersnatch.PointAffine, x fr.Element) []bandersnatch.PointAffine {

	if len(a) != len(b) {
		panic("slices not equal length")
	}

	result := make([]bandersnatch.PointAffine, len(a))
	for i := 0; i < len(a); i++ {
		var bx bandersnatch.PointAffine
		bx.ScalarMul(&b[i], &x)
		result[i].Add(&bx, &a[i])
	}
	return result
}

func generateChallenges(transcript *common.Transcript, proof *ipa.IPAProof) []fr.Element {

	challenges := make([]fr.Element, len(proof.L))
	for i := 0; i < len(proof.L); i++ {
		transcript.AppendPoint(&proof.L[i], "L")
		transcript.AppendPoint(&proof.R[i], "R")
		challenges[i] = transcript.ChallengeScalar("x")
	}
	return challenges
}

func domainToFr(in uint8) fr.Element {
	var x fr.Element
	x.SetUint64(uint64(in))
	return x
}

func arrayOfZeroes(n int) []fr.Element {
	r := make([]fr.Element, n)
	for i := 0; i < n; i++ {
		r[i] = fr.Zero()
	}
	return r[:]
}

func polynomialMul(a, b []fr.Element) []fr.Element {
	r := arrayOfZeroes(len(a) + len(b) - 1)
	for i := 0; i < len(a); i++ {
		for j := 0; j < len(b); j++ {
			// r[i+j] = fAdd(r[i+j], fMul(a[i], b[j]))
			var tmp fr.Element
			tmp.Mul(&a[i], &b[j])
			var res fr.Element
			res.Add(&r[i+j], &tmp)
			r[i+j] = res
		}
	}
	return r
}
